import sys
import argparse
import pysam

def autoLength(position,length):
    x1 = int(position[0])
    x2 = int(position[-1])
    dis = x2 - x1 - 1
    flank = int((length - dis) / 2)
    if (length - dis) % 2:
        left = x1 - flank + 1
        right = x2 + flank
    else:
        left = x1 - flank + 1 
        right = x2 + flank -1
    return [left - 1, right]


# def autoLength(position):
#     mid = (int(position[-1]) + int(position[0])) / 2 
#     left = int(mid) - 20
#     right = int(mid) + 20
#     return [left, right]

def replace_str(str,i,s):
    strlist = list(str)
    strlist[i] = s
    return ''.join(strlist)

def main():
    parser = argparse.ArgumentParser(description='')
    parser.add_argument("-f", "--fasta", dest="fasta", required=True, help="oligo fasta")
    parser.add_argument("-s", "--snp", dest="snp", required=True, help="snp information")
    parser.add_argument("-o", "--out", dest="out", required=True, help="output fasta")
    args = parser.parse_args()

    fastafile = pysam.FastaFile(args.fasta)

    f1 = open(args.out+".ref.fa", "w")
    f2 = open(args.out+".alt.fa", "w")

    with open (args.snp, "r") as f:
        for line in f:
            try:
                chrom, pos, mnvid, ref, alt, mnvrsid = line.rstrip().split("_")
            except ValueError:
                continue
            ref1 = ref.upper().split(",")
            alt1 = alt.upper().split(",")
            pos0 = autoLength(pos.split(","), 60)
            pos1 = pos.upper().split(",")
            # if len(ref) > 1 or len(alt) > 1:
            #     continue
            oligo_seq = fastafile.fetch(chrom, pos0[0], pos0[1]).upper()
            #print(oligo_seq)
            ref_seq = oligo_seq
            alt_seq = oligo_seq
            match = 1
            for i in range(len(ref1)):
                rpos = int(pos1[i]) - int(pos0[0]) - 1
                # print(pos[i])
                # print(pos0[0])
                if oligo_seq[rpos] != ref1[i] and oligo_seq[rpos] != alt1[i]:  
                    match = 0
                    break
                ref_seq = replace_str(ref_seq, rpos, ref1[i])
                alt_seq = replace_str(alt_seq, rpos, alt1[i])
            if not match:
                sys.stderr.write("{0} is not matched to the genome.\n".format(line.rstrip()))
                continue
            f1.write(">{0}_{1}_{2}_{3}_{4}_{5}\n".format(chrom, pos, mnvid, ref, alt, mnvrsid))
            f1.write("{0}\n".format(ref_seq))
            f2.write(">{0}_{1}_{2}_{3}_{4}_{5}\n".format(chrom, pos, mnvid, ref, alt, mnvrsid))
            f2.write("{0}\n".format(alt_seq))

    f1.close()
    f2.close()

if __name__ == "__main__":
    sys.exit(main())